/*
 * 2008 - 2009 Copyright (C) Ciprian Docan <docan@cac.rutgers.edu>
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation; either version 2 of the
 * License, or (at your option) any later version.
 */

#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <elf.h>

#include "pstrace.h"

extern int main( void );

#ifdef __x86_64__
typedef Elf64_Sym Elf_Sym; 
typedef Elf64_Shdr Elf_Shdr;
typedef Elf64_Ehdr Elf_Ehdr;
typedef Elf64_Addr Elf_Addr;

#define BP(bp) asm volatile ("mov %%rbp, %0" : "=r" (bp) : )

#else
typedef Elf32_Sym Elf_Sym;
typedef Elf32_Shdr Elf_Shdr;
typedef Elf32_Ehdr Elf_Ehdr;
typedef Elf32_Addr Elf_Addr;

#define BP(bp) asm volatile ("mov %%ebp, %0" : "=r" (bp) : )

#endif  /* __x86_64__ */

struct {
        char      *sym_name;
        int       sym_num;              /* Number of symbols. */
        Elf_Sym   *sym_tab;
} sym_struct;


static int pst_init = 0;

static int sym_compare( const void *a, const void *b )
{
        Elf_Sym *sym_a, *sym_b;

        sym_a = (Elf_Sym*) a;
        sym_b = (Elf_Sym*) b;

        if( sym_a->st_value < sym_b->st_value )
                return -1;
        else if( sym_a->st_value > sym_b->st_value )
                return 1;
        else return 0;
}


static void elf_read_symbol_table( int fd, Elf_Shdr* elfsh )
{
        sym_struct.sym_tab = (Elf_Sym*) malloc( elfsh->sh_size );
        if( !sym_struct.sym_tab )
                goto out; 

        if( lseek( fd, elfsh->sh_offset, SEEK_SET ) < 0 ) {
                perror( "seek failed" );
                goto out;
        }

        if( read( fd, sym_struct.sym_tab, elfsh->sh_size ) < 0 ) {
                perror( "read failed" );
                goto out;
        }

        sym_struct.sym_num = elfsh->sh_size / elfsh->sh_entsize;

        qsort( sym_struct.sym_tab, sym_struct.sym_num, 
               elfsh->sh_entsize, &sym_compare );

        printf( "there are %d symbols defined.\n", sym_struct.sym_num );
        /*
        for( i = 0; i < sym_struct.sym_num; i++ ) {
                printf( "%5d. 0x%8x %6d %s.\n", i,
                         sym_struct.sym_tab[i].st_value, 
                         sym_struct.sym_tab[i].st_size, 
                         sym_struct.sym_name + sym_struct.sym_tab[i].st_name );
                if( i == 93 )
                        printf( "offset = %d\n", sym_struct.sym_tab[i].st_name );
                * /
        }
        */
 out:
        ;
}


static void elf_read_strtab( int fd, Elf_Shdr *elfsh  )
{
        sym_struct.sym_name = (char*) malloc(elfsh->sh_size ) ;
        if( !sym_struct.sym_name )
                goto out;

        if( lseek( fd, elfsh->sh_offset, SEEK_SET ) < 0 ) {
                perror( "seek failed" );
                goto out;
        }

        if( read( fd, sym_struct.sym_name, elfsh->sh_size ) < 0 ) {
                perror( "read failed" );
                goto out;
        }

 out:
        ;
}

static void elf_read_section_header( int fd, Elf_Ehdr *elfhdr )
{
        Elf_Shdr *elfsh = 0;
        char     *buff = 0;
        int      size, i;


        size = elfhdr->e_shnum * elfhdr->e_shentsize;
        elfsh = (Elf_Shdr*) malloc( size );
        if( !elfhdr )
                goto out;

        if( lseek( fd, elfhdr->e_shoff, SEEK_SET ) < 0 ) {
                perror( "seek failed" );
                goto out;
        }

        if( read( fd, elfsh, size ) < 0 ) {
                perror( "read error" );
                goto out;
        }

        buff = (char*) malloc( elfsh[elfhdr->e_shstrndx].sh_size );
        if( !buff )
                goto out;

        if( lseek( fd, elfsh[elfhdr->e_shstrndx].sh_offset, SEEK_SET) < 0 ) {
                perror( "seek failed" );
                goto out;
        }

        if( read( fd, buff, elfsh[elfhdr->e_shstrndx].sh_size ) < 0 ) {
                perror( "read failed" );
                goto out;
        }

        /* for( i = 0; i < elfhdr->e_shnum; i++ ) { */
        for( i = elfhdr->e_shnum-1; i > -1; i-- ) {
                if( elfsh[i].sh_type == SHT_NULL )
                        printf( "%d. section name = NULL\n", i );
                else
                        printf( "%d. section name = %s\n", 
                                i, buff + elfsh[i].sh_name );

                if( elfsh[i].sh_type == SHT_SYMTAB ) {
                        elf_read_symbol_table( fd, elfsh+i );
                        break;
                }

                if( elfsh[i].sh_type == SHT_STRTAB ) {
                        printf( "found strtab ... " );
                        elf_read_strtab( fd, elfsh+i );
                }
        }

 out:
        if (elfsh)
                free(elfsh);

        if (buff)
                free(buff);
}


static void elf_resolve_symbol(Elf_Addr addr, char *buff)
{
        int left, right;
        int mid;
        char *sym_name;

        left = 0;
        right = sym_struct.sym_num - 1;

        while (left < right) {
                mid = (left + right) >> 1;
                if (sym_struct.sym_tab[mid].st_value < addr)
                        left = mid+1;
                else if (sym_struct.sym_tab[mid].st_value > addr)
                        right = mid-1;
                else    left = right = mid;
        }

        if (sym_struct.sym_tab[left].st_value > addr)
                left--;

        sym_name = sym_struct.sym_name + sym_struct.sym_tab[left].st_name;
        strcpy(buff, sym_name); // , strlen( sym_name) );
        /* printf( "0x%x, name = %s.\n", addr, sym_name ); */
}


static void my_function()
{
        char ch = 'a';

        printf( "character is %c\n", ch );
        pstrace();
}

/*
 * Initialize the library.
 */
void pstrace_init(void)
{
        char fname[80];
        int  fd;
        Elf_Ehdr elfhdr;

        if (pst_init) {
                printf( "pstrace already initialized.\n" );
                return;
        }

        pst_init = 1;

        sprintf(fname, "/proc/%d/exe", (int) getpid());
        fd = open(fname, O_RDONLY);
        if (fd < 0) {
                perror( "open failed" );
                return;
        }

        if (read(fd, &elfhdr, sizeof(Elf_Ehdr)) < 0) {
                perror("read header failed");
                goto out;
        }

        /*
        printf( "program entry:               0x%x\n"
                "program header offset:       %d\n"
                "section header offset:       %d\n"
                "program header entry size:   %d\n"
                "program header entries:      %d\n"
                "section header entry size:   %d\n"
                "section header entries:      %d\n"
                "section header string index: %d\n",
                elfhdr.e_entry,
                elfhdr.e_phoff,
                elfhdr.e_shoff,
                elfhdr.e_phentsize,
                elfhdr.e_phnum,
                elfhdr.e_shentsize,
                elfhdr.e_shnum,
                elfhdr.e_shstrndx );
        */
        elf_read_section_header(fd, &elfhdr);
 out:
        close( fd );
}

/*
  Function to print the stack call trace from the last function called
  down to the main routine.
*/
void pstrace(void)
{
        Elf_Addr *bp, entry_point;
        int i = 0, done = 0;
        char buff[256];

        if (!pst_init) {
                printf("pstrace not initialized, please call"
                        " pstrace_init() first");
                return;
        }

        entry_point = (Elf_Addr) &main;
        /* printf( "entry point at 0x%x.\n", entry_point ); */

	/*
	  The bottom two addresses on a stack frame contain the
	  address of the callee and the address of the next stack
	  frame; recover them in bp.
	*/
        BP(bp);
        do {
                /*
                if( bp[1] > entry_point ) 
                        done = 1;
                */
                elf_resolve_symbol(bp[1], buff);
                printf("[%d] 0x%x %s() \n", i, bp[1], buff);

                if (strcmp(buff, "main") == 0)
                        done = 1;

                bp = (Elf_Addr*) bp[0];
                i++;
        }
        while (!done);
}


void pstrace_finish()
{
        printf( "pstrace finish not implemented yet!\n" );
}


void pstrace_resolve_symbol(void *addr)
{
        char buff[256];

        elf_resolve_symbol((Elf_Addr) addr, buff);

        printf("0x%x %s\n", (Elf_Addr) addr, buff);
}


int a_main(void)
{
        printf( "test program, nothing special!\n" );

        pstrace_init();

        // elf_resolve_symbol( ((Elf_Addr) &main)  );
        elf_resolve_symbol( (Elf_Addr) &my_function, 0 );

        my_function();

        return 0;
}
